import os
import sys
import subprocess
import csv
import numpy as np
import matlab.engine
import copy
import time
import PX4MavCtrlV4 as PX4MavCtrl
import win32gui
import win32con
import pydirectinput

from attack_strategy import Ui_AttackStrategyConfig
from defense_strategy import Ui_DefenseStrategyConfig
from main_window import Ui_MainWindow
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QProgressBar, QLabel, QMessageBox, QDialog, QDesktopWidget)
from PyQt5.QtCore import *

scale_factor = 5
z_value = 30

class GetTrajWorker(QThread):
    finished_signal = pyqtSignal(np.ndarray)
    progress_signal = pyqtSignal(int)

    def __init__(self):
        super().__init__()

    def run(self):
        all_traj_xyz = self.get_traj()
        self.finished_signal.emit(all_traj_xyz)

    def get_traj(self):
        eng = matlab.engine.start_matlab()
        eng.cd('MatlabSrc\C4-main', nargout=0)
        y = np.array(eng.getTraj())
        eng.quit()
        all_traj_xyz = np.empty((125001, 12))
        traj_z = -z_value * np.ones_like(y[:, 0])
        col_sel = [12, 13, 16, 17, 20, 21, 24, 25]
        all_traj_xy = y[:, col_sel] * scale_factor
        for i in range(4):
            all_traj_xyz[:, 3 * i] = copy.deepcopy(all_traj_xy[:, 2 * i])
            all_traj_xyz[:, 3 * i + 1] = copy.deepcopy(all_traj_xy[:, 2 * i + 1])
            all_traj_xyz[:, 3 * i + 2] = copy.deepcopy(traj_z)
        return all_traj_xyz

class StartFlyWorker(QThread):
    signal_init = pyqtSignal(str)
    signal_fly = pyqtSignal(str)
    signal_track = pyqtSignal(str)
    signal_land = pyqtSignal(str)
    signal_no_traj = pyqtSignal(str)
    def __init__(self):
        self.all_traj_xyz = None
        super().__init__()
    def run(self):
        if self.all_traj_xyz is None:
            self.signal_no_traj.emit("请调用或等待matlab引擎运行完毕！")
        else:
            self.start_fly(self.all_traj_xyz)
    def start_fly(self,all_traj_xyz):
        vehicle_num = 4
        mav_list = []
        # Create MAV instance
        for i in range(vehicle_num):
            mav_list = mav_list + [PX4MavCtrl.PX4MavCtrler(1 + i)]
        time.sleep(2)
        # Start MAV loop with UDP mode: Udp_Simple
        for i in range(vehicle_num):
            mav_list[i].InitMavLoop(1)

        # Enter Offboard mode to start vehicle control
        time.sleep(2)
        for i in range(vehicle_num):
            mav_list[i].initOffboard()

        # Get the takeoff position of each vehicle to the UE4 Map
        # this can be adopted to obtain the global position of a vehicle in swarm simulation
        time.sleep(2)
        error2_ue4_map = []
        uav_global_pos = [[], [], [], []]
        for i in range(vehicle_num):
            mav = mav_list[i]
            error2_ue4_map = error2_ue4_map + [-np.array(
                [mav.uavGlobalPos[0] - mav.uavPosNED[0], mav.uavGlobalPos[1] - mav.uavPosNED[1],
                 mav.uavGlobalPos[2] - mav.uavPosNED[2]])]
            uav_global_pos[i].append(mav.pos)
        self.signal_init.emit('初始化传输链路完成')

        # fly to 15m high above its takeoff position
        for i in range(vehicle_num):
            mav_list[i].SendPosNED(0, 0, -15, 0)
            uav_global_pos[i].append(mav_list[i].pos)
        time.sleep(5)
        self.signal_fly.emit('起飞至高度 %f 米' % z_value)
        lastTime = time.time()
        startTime = time.time()
        timeInterval = 0.0004  # here is 0.0004s (2500Hz)
        self.signal_track.emit('开始目标追踪')
        while True:
            lastTime = lastTime + timeInterval
            sleepTime = lastTime - time.time()
            if sleepTime > 0:
                time.sleep(sleepTime)  # sleep until the desired clock
            else:
                lastTime = time.time()
            # The following code will be executed 2500Hz (0.0004s)

            t = time.time() - startTime
            t_step = int(t / 0.0004)
            if t_step >= 125000:  # end simulation when 50s
                break
            for j in range(vehicle_num):
                mav = mav_list[j]
                targetPosE = all_traj_xyz[t_step, 3 * j:3 * j + 3]
                # target position in vehicle takeoff frame
                targetPosE = targetPosE + error2_ue4_map[j]
                mav.SendPosNED(targetPosE[0], targetPosE[1], targetPosE[2], 0)
                uav_global_pos[j].append(mav.pos)

        for i in range(vehicle_num):
            mav_list[i].endOffboard()

        time.sleep(1)
        for i in range(vehicle_num):
            mav_list[i].stopRun()
        self.signal_track.emit('仿真结束，准备降落')

class AttackStrategyWindow(QWidget, Ui_AttackStrategyConfig):
    closing = pyqtSignal()
    def __init__(self, log_browser):
        super().__init__()
        # self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
        self.closing.connect(self.on_closing)
        self.setupUi(self)
        self.log_browser = log_browser
        self.submitBtn.clicked.connect(self.submit_btn_clicked)
        self.read_config_file()

    def read_config_file(self):
        file_path = os.path.join('Config', 'attack_strategy_config.csv')
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        self.lineEdit_DosP.setText(row.get('dos_period', ''))
                        self.doubleSpinBox_Ff.setValue(float(row.get('dos_freq', 0.1)))
                        self.doubleSpinBox_Impersonation.setValue(float(row.get('impersonation_signal', 0.01)))
                        self.doubleSpinBox_FDI.setValue(float(row.get('fdi_signal', 0.01)))
                        self.doubleSpinBox_AA_Signal.setValue(float(row.get('aa_signal', 0.01)))
                        self.doubleSpinBox_AA_Rate.setValue(float(row.get('aa_rate', 0.01)))
                self.append_log("INFO", "攻击策略配置文件读取成功")
            except Exception as e:
                self.append_log("ERROR", f"读取攻击策略配置文件时出错: {str(e)}")

    def closeEvent(self, event):
        self.closing.emit()
        event.accept()

    def on_closing(self):
        self.append_log("INFO", "防御策略配置窗口关闭")

    def submit_btn_clicked(self):
        dos_period = self.lineEdit_DosP.text()
        dos_freq = self.doubleSpinBox_Ff.value()
        impersonation_signal = self.doubleSpinBox_Impersonation.value()
        fdi_signal = self.doubleSpinBox_FDI.value()
        aa_signal = self.doubleSpinBox_AA_Signal.value()
        aa_rate = self.doubleSpinBox_AA_Rate.value()

        if dos_period == '':
            self.append_log("WARNING", "Dos攻击周期未定义，将使用默认值")
            dos_period = '1.11-5,11-15.3,21-24.9,31-34.9,41-44.9,51-54.9'
        if dos_freq == '':
            self.append_log("WARNING", "Dos攻击频率约束参数未定义，将使用默认值")
            dos_freq = 0.1
        if impersonation_signal == '':
            self.append_log("WARNING", "伪装攻击信号幅值未定义，将使用默认值")
            impersonation_signal = 0.01
        if fdi_signal == '':
            self.append_log("WARNING", "虚假注入攻击信号幅值未定义，将使用默认值")
            fdi_signal = 0.9
        if aa_signal == '':
            self.append_log("WARNING", "执行器攻击信号幅值未定义，将使用默认值")
            aa_signal = 0.01
        if aa_rate == '':
            self.append_log("WARNING", "执行器攻击变化速率上限未定义，将使用默认值")
            aa_rate = aa_signal * 13**0.5

        config_dir = 'Config'
        if not os.path.exists(config_dir):
            os.makedirs(config_dir)
        file_path = os.path.join(config_dir, 'attack_strategy_config.csv')
        try:
            with open(file_path, 'w', newline='', encoding='utf-8') as csvfile:
                fieldnames = ['dos_period', 'dos_freq', 'impersonation_signal', 'fdi_signal', 'aa_signal', 'aa_rate']
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerow({
                    'dos_period': dos_period,
                    'dos_freq': dos_freq,
                    'impersonation_signal': impersonation_signal,
                    'fdi_signal': fdi_signal,
                    'aa_signal': aa_signal,
                    'aa_rate': aa_signal * 13**0.5
                })
            self.append_log("SUCCESS", "攻击策略配置信息已保存到 Config/attack_strategy_config.csv")
            self.close()
        except Exception as e:
            self.append_log("ERROR", f"保存攻击策略配置信息时出错: {str(e)}")

    def append_log(self, level, message, timestamp=True):
        import time
        log_levels = {
            "SUCCESS": "<span style='color:#4CAF50; font-weight:bold;'>[成功]</span>",
            "INFO": "<span style='color:#2196F3;'>[信息]</span>",
            "WARNING": "<span style='color:#FFC107; font-weight:bold;'>[警告]</span>",
            "ERROR": "<span style='color:#F44336; font-weight:bold;'>[错误]</span>"
        }
        prefix = log_levels.get(level, "<span style='color:#666;'>[日志]</span>")
        time_str = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] " if timestamp else ""
        log_text = f"{time_str}{prefix} {message}\n"
        self.log_browser.append(log_text)
        self.log_browser.ensureCursorVisible()

class DefenseStrategyWindow(QWidget, Ui_DefenseStrategyConfig):
    closing = pyqtSignal()
    def __init__(self, log_browser):
        super().__init__()
        # self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
        self.closing.connect(self.on_closing)
        self.setupUi(self)
        self.log_browser = log_browser
        self.submitBtn.clicked.connect(self.submit_btn_clicked)
        self.read_config_file()

    def read_config_file(self):
        file_path = os.path.join('Config', 'defense_strategy_config.csv')
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        self.doubleSpinBox_EventTriggerSlope.setValue(float(row.get('event_trigger_slope', 0.6)))
                        self.doubleSpinBox_ZenoFreeInterval.setValue(float(row.get('zeno_free_interval', 0.205)))
                        self.doubleSpinBox_BackupTopologyTime.setValue(float(row.get('backup_topology_time', 2.538)))
                        self.lineEdit_ControllerGainK.setText(row.get('controller_gain_k', ''))
                        self.lineEdit_ControllerGainH.setText(row.get('controller_gain_h', ''))
                self.append_log("INFO", "防御策略配置文件读取成功")
            except Exception as e:
                self.append_log("ERROR", f"读取防御策略配置文件时出错: {str(e)}")

    def closeEvent(self, event):
        self.closing.emit()
        event.accept()

    def on_closing(self):
        self.append_log("INFO", "防御策略配置窗口关闭")

    def submit_btn_clicked(self):
        event_trigger_slope = self.doubleSpinBox_EventTriggerSlope.value()
        zeno_free_interval = self.doubleSpinBox_ZenoFreeInterval.value()
        backup_topology_time = self.doubleSpinBox_BackupTopologyTime.value()
        controller_gain_k = self.lineEdit_ControllerGainK.text()
        controller_gain_h = self.lineEdit_ControllerGainH.text()

        if event_trigger_slope == '':
            self.append_log("WARNING", "事件触发斜率g未定义,将采用默认值")
            event_trigger_slope = 0.6
        if zeno_free_interval == '':
            self.append_log("WARNING", "Zeno-free 时间间隔下界τ_Z未定义,将采用默认值")
            zeno_free_interval = 0.205
        if backup_topology_time == '':
            self.append_log("WARNING", "备份拓扑反应时间τ_R未定义,将采用默认值")
            backup_topology_time = 2.538
        if controller_gain_k == '':
            self.append_log("WARNING", "控制器增益矩阵K参数未定义,将采用默认值")
            controller_gain_k = '-0.4142,-0.6079'
        if controller_gain_h == '':
            self.append_log("WARNING", "控制器增益矩阵H参数未定义,将采用默认值")
            controller_gain_h = '0.9142,1.8079'

        config_dir = 'Config'
        if not os.path.exists(config_dir):
            os.makedirs(config_dir)
        file_path = os.path.join(config_dir, 'defense_strategy_config.csv')
        try:
            with open(file_path, 'w', newline='', encoding='utf-8') as csvfile:
                fieldnames = [
                    'event_trigger_slope', 'zeno_free_interval',
                    'backup_topology_time', 'controller_gain_k',
                    'controller_gain_h'
                ]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerow({
                    'event_trigger_slope': event_trigger_slope,
                    'zeno_free_interval': zeno_free_interval,
                    'backup_topology_time': backup_topology_time,
                    'controller_gain_k': controller_gain_k,
                    'controller_gain_h': controller_gain_h
                })
            self.append_log("SUCCESS", "防御策略配置信息已保存到 Config/defense_strategy_config.csv")
            self.close()
        except Exception as e:
            self.append_log("ERROR", f"保存防御策略配置信息时出错: {str(e)}")

    def append_log(self, level, message, timestamp=True):
        import time
        log_levels = {
            "SUCCESS": "<span style='color:#4CAF50; font-weight:bold;'>[成功]</span>",
            "INFO": "<span style='color:#2196F3;'>[信息]</span>",
            "WARNING": "<span style='color:#FFC107; font-weight:bold;'>[警告]</span>",
            "ERROR": "<span style='color:#F44336; font-weight:bold;'>[错误]</span>"
        }
        prefix = log_levels.get(level, "<span style='color:#666;'>[日志]</span>")
        time_str = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] " if timestamp else ""
        log_text = f"{time_str}{prefix} {message}\n"
        self.log_browser.append(log_text)
        self.log_browser.ensureCursorVisible()

class MainApplication(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        # self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
        self.init_connections()
        # 传递 logBrowser 给子窗口
        self.defense_window = DefenseStrategyWindow(self.logBrowser)
        self.attack_window = AttackStrategyWindow(self.logBrowser)
        # self.waiting_dialog = WaitingDialog()
        self.worker_get_traj = GetTrajWorker()
        self.worker_start_fly = StartFlyWorker()

    def init_connections(self):
        # 连接攻击配置界面按钮
        self.btnDefenseConfig.clicked.connect(self.open_defense_config)
        self.btnAttackConfig.clicked.connect(self.open_attack_config)
        self.btnRflysim.clicked.connect(self.start_rflysim)
        self.btnMatlab.clicked.connect(self.get_drone_trajectory)
        self.btnSimulate.clicked.connect(self.start_visual_simulation)
        self.btnExit.clicked.connect(self.exit_simulation)

    def open_defense_config(self):
        self.defense_window.read_config_file()
        self.defense_window.show()

    def open_attack_config(self):
        self.attack_window.read_config_file()
        self.attack_window.show()

    def start_rflysim(self):
        current_dir = os.path.dirname(os.path.abspath(__file__))
        bat_script_path = os.path.join(current_dir, "Bat", "1L3F.bat")

        try:
            process = subprocess.Popen(bat_script_path, shell=True)
            return_code = process.wait()

            if return_code == 0:
                self.append_log("SUCCESS", '批处理脚本启动成功，请等待无人机初始化完毕')
            else:
                self.append_log("ERROR", '批处理脚本执行出错')
        except FileNotFoundError:
            self.append_log("ERROR", f"批处理脚本 {bat_script_path} 不存在")
        except Exception as e:
            self.append_log("ERROR", f'其他错误: {str(e)}')

    def get_drone_trajectory(self):
        self.worker_get_traj.finished_signal.connect(self.save_traj)
        # self.waiting_dialog.show()
        self.append_log("INFO", "正在获取无人机轨迹，请等待完成获取的消息提醒")
        self.worker_get_traj.start()

    def save_traj(self, all_traj_xyz):
        self.worker_start_fly.all_traj_xyz = all_traj_xyz
        # self.waiting_dialog.close()
        self.append_log("SUCCESS", "已经获取无人机轨迹")

    def start_visual_simulation(self):
        self.worker_start_fly.signal_init.connect(lambda msg: self.append_log("SUCCESS", msg))
        self.worker_start_fly.signal_fly.connect(lambda msg: self.append_log("INFO", msg))
        self.worker_start_fly.signal_track.connect(lambda msg: self.append_log("INFO", msg))
        self.worker_start_fly.signal_land.connect(lambda msg: self.append_log("SUCCESS", msg))
        self.worker_start_fly.signal_no_traj.connect(lambda msg: self.append_log("WARNING", msg))

        rfly_window = self.find_rfly_window()
        if rfly_window:
            if win32gui.IsIconic(rfly_window):
                win32gui.ShowWindow(rfly_window, win32con.SW_RESTORE)
            else:
                win32gui.SetForegroundWindow(rfly_window)
        else:
            self.append_log("WARNING", "仿真平台未启动！请先初始化Rflysim仿真平台")
            return
        try:
            time.sleep(2)
            pydirectinput.press('s')
            time.sleep(0.5)
            pydirectinput.typewrite('t5', interval=0.1)
            self.append_log("INFO", "即将开始可视化仿真")
            self.worker_start_fly.start()
        except Exception as e:
            self.append_log("ERROR", f"启动可视化仿真时出错: {str(e)}")


    def find_windows(self):
        hwnd_list = []

        def enum_callback(hwnd, _):
            if win32gui.GetClassName(hwnd) in ['ConsoleWindowClass', 'UnrealWindow']:
                hwnd_list.append(hwnd)

        win32gui.EnumWindows(enum_callback, None)
        return hwnd_list

    def find_cmd_window(self):
        cmd_windows = self.find_windows()
        for window in cmd_windows:
            if 'C:\\Windows\\System32\\cmd.exe' in win32gui.GetWindowText(window):
                # self.get_window_text(window)
                return window
        return None

    def find_rfly_window(self):
        cmd_windows = self.find_windows()
        for window in cmd_windows:
            if 'Rfly' in win32gui.GetWindowText(window):
                return window
        return None

    def get_window_text(self,hwnd):
        """获取窗口文本内容（适用于可编辑文本框或标题）"""
        text_length = win32gui.GetWindowTextLength(hwnd)
        text_buffer = win32gui.PyMakeBuffer(text_length + 1)
        win32gui.GetWindowText(hwnd, text_buffer, text_length + 1)
        print(text_buffer[:-1].decode('gbk', errors='ignore'))

    def exit_simulation(self):
        cmd_window = self.find_cmd_window()
        if cmd_window:
            if win32gui.IsIconic(cmd_window):
                win32gui.ShowWindow(cmd_window, win32con.SW_RESTORE)
            else:
                win32gui.SetForegroundWindow(cmd_window)

            try:
                # time.sleep(2)
                pydirectinput.press('enter')
                time.sleep(4)
                if self.find_rfly_window() is not None:
                    self.append_log("WARNING", f"退出仿真时出错,请手动关闭")
                else:
                    self.append_log("SUCCESS", "Rflysim平台已退出！")
            except Exception as e:
                self.append_log("ERROR", f"退出仿真时出错: {str(e)}")
        else:
            self.append_log("WARNING", '仿真未启动！')

    def append_log(self, level, message, timestamp=True):
        log_levels = {
            "SUCCESS": "<span style='color:#4CAF50; font-weight:bold;'>[成功]</span>",
            "INFO": "<span style='color:#2196F3;'>[信息]</span>",
            "WARNING": "<span style='color:#FFC107; font-weight:bold;'>[警告]</span>",
            "ERROR": "<span style='color:#F44336; font-weight:bold;'>[错误]</span>"
        }
        prefix = log_levels.get(level, "<span style='color:#666;'>[日志]</span>")
        time_str = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] " if timestamp else ""
        log_text = f"{time_str}{prefix} {message}\n"
        self.logBrowser.append(log_text)
        self.logBrowser.ensureCursorVisible()